import networkx as nx
import matplotlib.pyplot as plt
G = nx.DiGraph()
G.add_edge('A','B',capacity = 4)
G.add_edge('A','C',capacity = 5)
G.add_edge('B','C',capacity = 2)
G.add_edge('B','D',capacity = 6)
G.add_edge('C','D',capacity = 1)
G.add_edge('C','E',capacity = 3)
G.add_edge('D','E',capacity = 8)
G.add_nodes_from(['A','B','C','D','E'])
nx.draw(G,with_labels=True)
a = G['A']['B']['capacity']
b = G['A'].items()
#print(b)

def initialize_residual_network(G):
    residual_network = G.copy()
    for u,v,attr in residual_network.edges(data = True):
        attr['residual_capacity'] = attr['capacity']
    return residual_network

def find_augmenting_path(residual_network,source,sink):
    path = []
    visited = {node: False for node in residual_network.nodes()}

    def dfs(node ,current_path):
        if node == sink:
            path.extend(current_path)
            return True
        visited[node] = True

        for neighbor, attr in residual_network[node].items():
            if not visited[neighbor] and attr['residual_capacity']>0:
                found_sink = dfs(neighbor,current_path+[(node,neighbor)])
                if found_sink:
                    return True

        return False
    if dfs(source,[]):
        return path
    return None

def calculate_flow(path, residual_network):
    flow = float('inf')  # 设置流量为无穷大
    for u, v in path:
        flow = min(flow, residual_network[u][v]['residual_capacity'])  # 取最小残余容量作为流量
    return flow


def update_residual_network(residual_network,augmenting_path,flow):
    for u,v in augmenting_path:
        residual_network[u][v]['residual_capacity'] -=flow

        if residual_network.has_edge(v,u):
            residual_network[v][u]['residual_capacity'] += flow
        else:
            residual_network.add_edge(v,u,residual_capacity = flow)




def ford_fulkerson(G,source,sink):
    residual_network = initialize_residual_network(G)
    max_flow = 0

    while True:
        agumenting_path = find_augmenting_path(residual_network,source,sink)
        if not agumenting_path:
            break
        flow = calculate_flow(agumenting_path,residual_network)
        max_flow += flow
        update_residual_network(residual_network,agumenting_path,flow)
    return max_flow,residual_network



def find_min_cut(residual_network,source):
    reachable = bfs_reachable(residual_network,source)
    min_cut = []
    for u in reachable:
        for v,attr in residual_network[u].items():
            if attr['residual_capacity'] == 0 and v not in reachable:
                min_cut.append((u,v))
    return min_cut

def bfs_reachable(residual_network,source):
    visited = {node: False for node in residual_network.nodes()}
    visited[source] = True
    queue = [source]

    while queue:
        node = queue.pop(0)
        for neighbor, attr in residual_network[node].items():
            if not visited[neighbor] and attr['residual_capacity'] > 0:
                visited[neighbor] = True
                queue.append(neighbor)
    return [node for node,is_reachable in visited.items() if is_reachable]


def Cal_widening_oneRoad_max_flow(G,u,v):
    G_copy = G.copy()
    if G_copy.has_edge(u,v):
        G_copy[u][v]['capacity'] += 1
    widening_oneRoad_max_flow = ford_fulkerson(G_copy, 'A', 'E')[0]
    return widening_oneRoad_max_flow

def Cal_widening_OneRoad_max_flows(G,min_cut):
    max_flows = []
    for u,v in min_cut:
        if G.has_edge(u,v):
            widening_flow = Cal_widening_oneRoad_max_flow(G,u,v)
            max_flows.append((u,v,widening_flow))
    max_flows.sort(key = lambda x:x[2],reverse=True)
    return max_flows



# 算原来流网络的最大流
Old_max_flow,residual_network = ford_fulkerson(G,'A','E')
print(Old_max_flow)
# 找出最小割
mincut = find_min_cut(residual_network,'A')
print(mincut)
# 算出扩容每一条边对应的最大流
max_flows = Cal_widening_OneRoad_max_flows(G,mincut)
print(max_flows)
# # 验算扩容后的有向图
# def print_graph_edges_with_capacity(G):
#     print("Edges and their capacities after expansion:")
#     for u, v, attr in G.edges(data=True):
#         print(f"Edge ({u} -> {v}) with capacity {attr['capacity']}")
#
# # 扩容后查看所有边的容量
# print_graph_edges_with_capacity(G)
#
# nx.draw(G,with_labels=True)
# # 扩容后求最大流
# Widening_Max_flow,residual_network = ford_fulkerson(G,'A','E')
# print(Widening_Max_flow)





